#!/usr/bin/python
# -*-coding:utf-8-*-
import os
import numpy as np
import pandas as pd

### 底层读取数据的依赖（不提供）
from zbc_factor_lib.base.factors_library_base import NewRQFactorLib as DataReader

from dir_info import label_data_dir

data_reader = DataReader()
# data_reader.show_basic_library_db()

# TODO - given date then map to int index
def trade_date_map_int(date):
    trade_date_data_map_df = pd.read_hdf(os.path.join(label_data_dir, 'trade_date_data_map_df.h5'))

    idx = trade_date_data_map_df.loc[pd.to_datetime(date), 'label']

    return idx

def trade_date_map(df, key='date', reverse=False, is_financial_date=False):
    df = df.copy()

    if is_financial_date:
        trade_date_data_map_df = pd.read_hdf(os.path.join(label_data_dir, 'financial_date_data_map_df.h5'))
    else:
        trade_date_data_map_df = pd.read_hdf(os.path.join(label_data_dir, 'trade_date_data_map_df.h5'))

    if reverse:
        trade_date_data_map_df['trade_date'] = trade_date_data_map_df.index

        trade_date_data_map_df = trade_date_data_map_df.set_index('label')

        df[key] = trade_date_data_map_df.loc[df[key], 'trade_date'].values

        df[key] = pd.to_datetime(df[key])
    else:
        df[key] = trade_date_data_map_df.loc[pd.to_datetime(df[key]), 'label'].values

        df[key] = df[key].astype(np.uint16)

    return df


def stock_code_map(df, key='stock_code', reverse=False):
    df = df.copy()

    stock_code_map = pd.read_hdf(os.path.join(label_data_dir, 'stock_code_map.h5'))

    if reverse:
        stock_code_map['stock_code'] = stock_code_map.index

        stock_code_map = stock_code_map.set_index('label')

        df[key] = stock_code_map.loc[df[key], 'stock_code'].values
    else:
        df[key] = stock_code_map.loc[df[key], 'label'].values

        df[key] = df[key].astype(np.uint16)

    return df


def ci1_code_map(df, key='ci1_code', reverse=False):
    df = df.copy()

    ci1_code_map = pd.read_hdf(os.path.join(label_data_dir, 'ci1_code_map.h5'))

    if reverse:
        ci1_code_map['index_code'] = ci1_code_map.index

        ci1_code_map = ci1_code_map.set_index('label')

        df[key] = ci1_code_map.loc[df[key], 'index_code'].values
    else:
        df[key] = ci1_code_map.loc[df[key], 'label'].values

        df[key] = df[key].astype(np.uint8)

    return df


#####################################
def main():
    financial_date_data = ['%s-03-31' % year for year in range(2007, 2030)] + \
                          ['%s-06-30' % year for year in range(2007, 2030)] + \
                          ['%s-09-30' % year for year in range(2007, 2030)] + \
                          ['%s-12-31' % year for year in range(2007, 2030)]

    financial_date_data = sorted(financial_date_data)

    financial_date_data = pd.DatetimeIndex(financial_date_data)

    financial_date_data = pd.DataFrame({'label': range(financial_date_data.shape[0])}, index=financial_date_data)

    financial_date_data['label'] = financial_date_data['label'].astype(np.uint16)

    financial_date_data.to_hdf(os.path.join(label_data_dir, 'financial_date_data_map_df.h5'),
                                  key='financial_date_data_map_df')

    ##
    trade_date_data = data_reader.read_basic_data_table('processed_trade_date_data')

    trade_date_data = pd.DatetimeIndex(trade_date_data['trade_date'].unique())

    trade_date_data_map_df = pd.DataFrame({'label': range(trade_date_data.shape[0])}, index=trade_date_data)

    trade_date_data_map_df['label'] = trade_date_data_map_df['label'].astype(np.uint16)

    trade_date_data_map_df.to_hdf(os.path.join(label_data_dir, 'trade_date_data_map_df.h5'),
                                  key='trade_date_data_map_df')

    ##
    processed_stock_list_delist_info_data = data_reader.read_basic_data_table('processed_stock_list_delist_info_data')

    processed_stock_list_delist_info_data = processed_stock_list_delist_info_data.set_index('code')

    processed_stock_list_delist_info_data = processed_stock_list_delist_info_data.sort_index()

    processed_stock_list_delist_info_data['label'] = range(processed_stock_list_delist_info_data.shape[0])

    stock_code_map = processed_stock_list_delist_info_data[['label']]

    stock_code_map['label'] = stock_code_map['label'].astype(np.uint16)


    stock_code_map.to_hdf(os.path.join(label_data_dir, 'stock_code_map.h5'), key='stock_code_map')

    ## ci1_code map
    processed_ci1_industry_index_info_data = data_reader.read_basic_data_table('processed_ci1_industry_index_info_data')

    processed_ci1_industry_index_info_data = processed_ci1_industry_index_info_data.set_index('index_code')

    processed_ci1_industry_index_info_data = processed_ci1_industry_index_info_data.sort_index()

    processed_ci1_industry_index_info_data['label'] = range(processed_ci1_industry_index_info_data.shape[0])

    ci1_code_map = processed_ci1_industry_index_info_data[['label']]

    ci1_code_map['label'] = ci1_code_map['label'].astype(np.uint8)

    ci1_code_map.to_hdf(os.path.join(label_data_dir, 'ci1_code_map.h5'), key='ci1_code_map')


if __name__ == '__main__':
    main()

